import re,os,sys,numpy as np
from scipy import stats
import matplotlib.pyplot as pyplot
import warnings
warnings.simplefilter('ignore', DeprecationWarning)

def plotcompare(tmin,tmax,nbins,sdevs):
    # bin the data so we can plot it
    bincenters = tmin + (tmax-tmin)*.01*(np.array(range(0,nbins))+.5)
    #print bincenters
    binvals = np.zeros(nbins)

    for point in sdevs:
        difs = np.abs(point - bincenters)
        bin = np.argmin(difs)
        binvals[bin] += 1
    binvals /= float(len(sdevs))

    distbinvals = np.zeros(nbins)
    distbincenters = tmin + (tmax-tmin)*.01*(np.array(range(0,nbins))+.5)
    data = stats.gamma.rvs(alpha,loc=loc,scale=beta,size=nbins)
    for point in data:
        difs = np.abs(point - bincenters)
        bin = np.argmin(difs)
        distbinvals[bin] += 1

    distbinvals /= float(len(data))
    #plotbinvals = np.cumsum(binvals)
    #binvals = plotbinvals

    #numargs = stats.gamma.numargs
    #[ a ] = [alpha,] * numargs
    #rv = stats.gamma(a,loc=loc,scale=beta)
    #plotpoints = tmin + (tmax/1000.-tmin)*.01*(np.array(range(0,nbins))+.001)

    #pyplot.plot(bincenters, rv.cdf(bincenters),c="r")
    #pyplot.scatter(plotpoints, rv.pdf(plotpoints),c="r")

    pyplot.scatter(bincenters,binvals)
    pyplot.scatter(distbincenters,distbinvals,c="r")
    #pyplot.scatter(bincenters[:10],binvals[:10])
    #pyplot.plot(bincenters,stats.gamma.pdf(bincenters,loc=loc,scale=beta),c="r")
    pyplot.show()

    #print vals[np.isfinite(vals)].mean()

def ddists(basetracksdir,fname):
    posdicts = dict()
    for trackfname in sorted(os.listdir(basetracksdir+"/"+fname)):
        groupings = trackfname.split(".")
        trackname = groupings[0].lower()
        if trackname.startswith("person"):
            trackhandle = open(basetracksdir+"/"+fname+"/"+trackfname,"r")
            coords = trackhandle.readlines()
            xs = np.zeros(len(coords))
            zs = np.zeros(len(coords))
            for i,vec in enumerate(coords):
                vecels = re.split("\s+",vec.strip())
                xs[i] = vecels[0] if vecels[3].startswith("1") else float("nan")
                zs[i] = vecels[2] if vecels[3].startswith("1") else float("nan")
                posdicts[trackname] = (xs,zs)
    alldists = []
    for track0 in sorted(posdicts.keys()):
        for track1 in sorted(posdicts.keys()):
            if track1 <= track0:
                continue
            x0 = posdicts[track0][0]
            x1 = posdicts[track1][0]
            z0 = posdicts[track0][1]
            z1 = posdicts[track1][1]
            xd = x0 - x1
            zd = z0 - z1
            dists = np.sqrt(xd*xd+zd*zd)
            dists0 = np.concatenate((np.zeros(1),dists))
            dists0[0] = float("nan")
            dists1 = np.concatenate((dists,np.zeros(1)))
            dists1[len(dists1)-1] = float("nan")
            ddists = dists1-dists0
            ddists = ddists[1:-1]
            alldists.append(ddists)
    return alldists

#runs from msave dir
basetracksdir = "./data/"
posdicts = dict()

if len(sys.argv) > 1:
    basetracksdir += sys.argv[1]
    alldists = []
    alldists.append(ddists(basetracksdir))
    alldists = np.array(alldists)
    
else:
    alldists = []
    for subdir in sorted(os.listdir(basetracksdir)):
        curdists = ddists(basetracksdir,subdir)
        #print curdists
        for subarray in curdists:
            alldists.append(subarray)
        #print alldists

#print alldists
maxlen = max(map(lambda x: len(x), alldists))
wsize = 15
sframe = 0
ssize = 1
means = []
sdevs = []

while sframe+wsize <= maxlen:
    window = np.array([])
    for dlist in alldists:
        if sframe+wsize > len(dlist):
            continue
        window = np.concatenate((window,np.array(dlist[sframe:sframe+wsize])))
    # now we don't care about temporal ordering so we can just throw out nan frames
    window = window[np.isfinite(window)]
    means.append(np.mean(window))
    std = np.std(window)
    sdevs.append(1./std/std) #want the precision, not the sdev
    sframe += ssize

means = np.array(means)
sdevs = np.array(sdevs)
tmean = np.mean(means)
tsdev = np.mean(sdevs)

print tmean,tsdev
alpha,loc,beta = stats.gamma.fit(sdevs)
beta = 1./beta #parameters are actually k,theta
print "alpha",alpha
print "beta",beta
outfile = open('distparams','w')
outfile.write(str(alpha)+"\n"+str(beta)+"\n")

tmax = np.max(sdevs)
tmin = np.min(sdevs)
print "min",tmin
print "max",tmax

#plotcompare(tmin,tmax,1000,sdevs)
